import subprocess
import os

def split_video_by_pipe(video_path, frame_segments, width, height, output_dir):
    """
    使用 FFmpeg 的管道方式，将视频拆分成指定的帧切片，输出为多个视频文件。

    参数：
    - video_path: 输入视频路径
    - frame_segments: 切片帧段，例如：[0, M1, M2, ..., Mlast]
    - width, height: 视频分辨率，用于计算帧大小
    - output_dir: 输出文件目录
    """
    # 视频帧格式假设 RGB（每像素 3 字节）
    bytes_per_pixel = 3
    frame_size = width * height * bytes_per_pixel  # 单帧的字节大小

    # 创建输出子目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 打开视频解码管道，用于帧提取
    ffmpeg_extract = [
        "ffmpeg", "-i", video_path, "-vf", "fps=30", "-f", "image2pipe", "-vcodec", "rawvideo", "-"
    ]
    process_extract = subprocess.Popen(ffmpeg_extract, stdout=subprocess.PIPE, bufsize=frame_size * 1)

    # 创建多个帧分段管道
    slice_pipes = []
    for i, _ in enumerate(frame_segments[:-1]):
        output_file = os.path.join(output_dir, f"slice_{i+1}.mp4")
        ffmpeg_encode = [
            "ffmpeg", "-framerate", "30", "-f", "rawvideo", "-vcodec", "rawvideo", "-s", f"{width}x{height}", "-pix_fmt", "rgb24", "-i", "-",
            "-c:v", "libx264", "-preset", "slow", "-crf", "18", "-pix_fmt", "yuv420p", output_file
        ]
        slice_pipes.append(subprocess.Popen(ffmpeg_encode, stdin=subprocess.PIPE, bufsize=frame_size * 1))

    # 读取帧并分配到对应的切片管道
    current_frame = 0
    while True:
        frame_data = process_extract.stdout.read(frame_size)
        if len(frame_data) < frame_size:
            break  # 读取结束

        # 找到当前帧所属的切片段
        for i, (start_frame, end_frame) in enumerate(zip(frame_segments[:-1], frame_segments[1:])):
            if start_frame <= current_frame < end_frame:
                print(i)
                slice_pipes[i].stdin.write(frame_data)
                break
        
        current_frame += 1

    # 关闭所有管道
    process_extract.stdout.close()
    process_extract.wait()
    for pipe in slice_pipes:
        pipe.stdin.close()
        pipe.wait()

    print(f"视频分段完成，切片存储在：{output_dir}")



# 示例使用
video_path = "E:/抖音-AI风景/Test/demo.mp4"
frame_segments = [0, 60, 69, 305]
fps = 30
width, height = 1080, 1440  
output_dir = "output"

split_video_by_pipe(video_path, frame_segments, width,height, output_dir)